package com.google.code.twiddling.core.util;

import java.io.BufferedInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileFilter;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.annotation.Annotation;
import java.lang.reflect.AnnotatedElement;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

import org.objectweb.asm.AnnotationVisitor;
import org.objectweb.asm.Attribute;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.FieldVisitor;
import org.objectweb.asm.MethodVisitor;

/**
 * @author <a href="mailto:yy.young@gmail.com">Young Yang</a>
 * @create Dec 16, 2008 2:07:04 PM
 */
public class ASMUtils {

    /**
     * parse annonation in the scope of urlCL, and find out the annotated class
     * @param urlCL
     * @param annotation
     * @return
     */
    public static Class[] parseAnnoation(URLClassLoader urlCL, Class<? extends Annotation> annotation) {
        ASMParser asmParser = new ASMParser(urlCL);
        asmParser.parseClasses(urlCL.getURLs());
        return asmParser.findClassAnnotatedWith(annotation);
    }


    private static class ASMParser {

        private URLClassLoader cl = null;

        /**
         * AnnotationName=>[ClassInfo]
         */
        private final Map<String, List<ClassInfo>> annotated = new HashMap<String, List<ClassInfo>>();

        public ASMParser(URLClassLoader cl) {
            this.cl = cl;
        }

        /**
         *
         * @param annotation annotation
         */
        public Class[] findClassAnnotatedWith(Class<? extends Annotation> annotation) {
            List<Class> classes = new ArrayList<Class>();
            List<ClassInfo> infos = getAnnotationInfos(annotation.getName());
            for (ClassInfo classInfo : infos) {
                try {
                    Class clazz = classInfo.get();
                      // double check via proper reflection
                    if (clazz.isAnnotationPresent(annotation)) {
                        classes.add(clazz);
                    }
                }
                catch (ClassNotFoundException e) {
                    throw new RuntimeException("Exception occupied while findClassAnnotatedWith " + annotation, e);
                }
            }
            return classes.toArray(new Class[classes.size()]);
        }

        public void parseClasses(URL[] urls) {
            List<URL> appURLs = new ArrayList<URL>();
            for (URL url : urls) {
                URLClassLoader urlClassLoader = new URLClassLoader(new URL[]{url}, null);
                URL testURL = urlClassLoader.findResource(Object.class.getName().replace(".", "/") + ".class");
                URL charsetURL = urlClassLoader.findResource("sun/nio/cs/ext/GBK.class"); // charset.jar
                if (testURL == null && charsetURL == null) {
                    appURLs.add(url);
                }
            }
            Map<String, byte[]> classBytesArray = new HashMap<String, byte[]>();
            for (URL url : appURLs) {
                try {
                    classBytesArray.putAll(getClassBytesMap(url));
                }
                catch (IOException e) {
                    throw new RuntimeException("Failed to get Class names from url: " + url.toString());
                }
            }
            for (Map.Entry<String, byte[]> entry : classBytesArray.entrySet()) {
                readClass(entry.getKey(), entry.getValue());
            }

        }


// ----------- ASM bytecode reader ------------

        private void readClass(String className, byte[] classBytes) {
            try {
                ClassReader classReader = new ClassReader(classBytes);
                classReader.accept(new ClassInfoBuildingVisitor(), ClassReader.SKIP_DEBUG);
            }
            catch (RuntimeException e) {
                throw new RuntimeException("Error to read class: " + className, e);
            }
        }

        private List<ClassInfo> getAnnotationInfos(String name) {
            List<ClassInfo> infos = annotated.get(name);
            if (infos == null) {
                infos = new ArrayList<ClassInfo>();
                annotated.put(name, infos);
            }
            return infos;
        }

        class ClassInfoBuildingVisitor implements ClassVisitor, AnnotationVisitor {
            private AnnotatableInfo info;

            public ClassInfoBuildingVisitor() {
            }

            public ClassInfoBuildingVisitor(AnnotatableInfo info) {
                this.info = info;
            }

            public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
                if (!name.endsWith("package-info")) {
                    ClassInfo classInfo = new ClassInfo(javaName(name), javaName(superName));

                    for (String interfce : interfaces) {
                        classInfo.getInterfaces().add(javaName(interfce));
                    }
                    info = classInfo;
                }
            }

            public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
                AnnotationInfo annotationInfo = new AnnotationInfo(desc);
                info.getAnnotations().add(annotationInfo);
                if (info instanceof ClassInfo) {
                    getAnnotationInfos(annotationInfo.getName()).add((ClassInfo) info);
                }
                return new ClassInfoBuildingVisitor(annotationInfo);
            }

            private String javaName(String name) {
                return (name == null) ? null : name.replace('/', '.');
            }

            //--------- not need to implement
            public void visitAttribute(Attribute attribute) {
            }

            public void visitEnd() {
            }

            public FieldVisitor visitField(int i, String string, String string1, String string2, Object object) {
                return null;
            }

            public void visitInnerClass(String string, String string1, String string2, int i) {
            }

            public MethodVisitor visitMethod(int i, String string, String string1, String string2, String[] strings) {
                return null;
            }

            public void visitOuterClass(String string, String string1, String string2) {
            }

            public void visitSource(String string, String string1) {
            }

            public void visit(String string, Object object) {
            }

            public AnnotationVisitor visitAnnotation(String string, String string1) {
                return null;
            }

            public AnnotationVisitor visitArray(String string) {
                return null;
            }

            public void visitEnum(String string, String string1, String string2) {
            }
        }

        abstract class AnnotatableInfo {
            private final List<AnnotationInfo> annotations = new ArrayList<AnnotationInfo>();

            public AnnotatableInfo() {
            }

            public AnnotatableInfo(AnnotatedElement element) {
                for (Annotation annotation : element.getAnnotations()) {
                    annotations.add(new AnnotationInfo(annotation.annotationType().getName()));
                }
            }

            public List<AnnotationInfo> getAnnotations() {
                return annotations;
            }

            public abstract String getName();
        }

        class ClassInfo extends AnnotatableInfo {
            private final String name;
            private final String superType;
            private final List<String> interfaces = new ArrayList<String>();
            private Class<?> clazz;
            private ClassNotFoundException notFound;

            public ClassInfo(Class clazz) {
                super(clazz);
                this.clazz = clazz;
                this.name = clazz.getName();
                Class superclass = clazz.getSuperclass();
                this.superType = superclass != null ? superclass.getName() : null;
            }

            public ClassInfo(String name, String superType) {
                this.name = name;
                this.superType = superType;
            }

            public String getPackageName() {
                return name.substring(name.lastIndexOf(".") + 1, name.length());
            }

            public List<String> getInterfaces() {
                return interfaces;
            }

            public String getName() {
                return name;
            }

            public String getSuperType() {
                return superType;
            }

            public Class get() throws ClassNotFoundException {
                if (clazz != null) return clazz;
                if (notFound != null) throw notFound;
                try {
                    this.clazz = ASMParser.this.cl.loadClass(name);
                }
                catch (ClassNotFoundException e) {
                    notFound = e;
                    throw e;
                }
                return clazz;
            }

            public String toString() {
                return name;
            }
        }

        class AnnotationInfo extends AnnotatableInfo {
            private final String name;

            public AnnotationInfo(Annotation annotation) {
                this(annotation.getClass().getName());
            }

            public AnnotationInfo(Class<? extends Annotation> annotation) {
                this.name = annotation.getName().intern();
            }

            public AnnotationInfo(String name) {
                name = name.replaceAll("^L|;$", "");
                name = name.replace('/', '.');
                this.name = name.intern();
            }

            public String getName() {
                return name;
            }

            public String toString() {
                return getName();
            }
        }

    }

    public static Map<String, byte[]> getClassBytesMap(URL location) throws IOException {
        File file = toFile(location);
        return getClassBytesMap(file);
    }

    public static Map<String, byte[]> getClassBytesMap(File file) throws IOException {
        Map<String, byte[]> contents = new HashMap<String, byte[]>();
        if (file.isDirectory()) {
            List<File> files = listFiles(file,
                    new FileFilter() {
                        public boolean accept(File pathname) {
                            return pathname.getName().endsWith(".class") && !pathname.getName().equals("package-info.class");
                        }
                    }
            );
            for (File _file : files) {
                String className = _file.getPath().substring(file.getPath().length() + 1);
                className = className.replaceFirst(".class$", "");
                className = className.replace(File.separatorChar, '.');
                byte[] content = toByteArray(new FileInputStream(_file));
                contents.put(className,content);
            }
        }
        else {
            JarFile jarFile = new JarFile(file);
            Enumeration<JarEntry> enu = jarFile.entries();
            try {
                while (enu.hasMoreElements()) {
                    JarEntry entry = enu.nextElement();
                    if (entry.isDirectory() || !entry.getName().endsWith(".class") || entry.getName().endsWith("package-info.class")) {
                        continue;
                    }
                    String className = entry.getName();
                    className = className.replaceFirst(".class$", "");
                    className = className.replace('/', '.');
                    InputStream in = null;
                    try {
                        in = new BufferedInputStream(jarFile.getInputStream(entry));
                        byte[] content = toByteArray(in);
                        contents.put(className, content);
                    }
                    finally {
                        if(in != null) {
                            in.close();
                        }
                    }
                }
            }
            finally {
                jarFile.close();
            }
        }
        return Collections.unmodifiableMap(contents);
    }

    private static File toFile(URL url) {
        if (!url.getProtocol().equals("file")) {
            return null;
        }
        else {
            String filename = url.getFile().replace('/', File.separatorChar);
            return new File(filename);
        }
    }

    private static List<File> listFiles(File dir, FileFilter filter) {
        List<File> files = new ArrayList<File>();
        if (!dir.exists() || dir.isFile())
            return files;
        listFiles(files, dir, filter);
        return files;
    }

    private static void listFiles(List filesList, File dir, FileFilter filter) {
        File[] files = dir.listFiles(filter);
        List temp = Arrays.asList(files);
        Collections.sort(temp);
        filesList.addAll(temp);

        File[] subDirs = dir.listFiles(new FileFilter() {
            public boolean accept(File pathname) {
                return pathname.isDirectory();
            }
        });
        for (int i = 0; i < subDirs.length; i++) {
            listFiles(filesList, subDirs[i], filter);
        }
    }

    private static byte[] toByteArray(InputStream input) throws IOException {
        ByteArrayOutputStream output = new ByteArrayOutputStream();
        copy(input, output);
        return output.toByteArray();
    }

    private static int copy(InputStream input, OutputStream output)
            throws IOException {
        byte[] buffer = new byte[1024];
        int count = 0;
        int n = 0;
        while (-1 != (n = input.read(buffer))) {
            output.write(buffer, 0, n);
            count += n;
        }
        return count;
    }
    
}
